import copy
from matplotlib import pyplot as plt
from matplotlib import animation

# 训练数据集
training_set = [[(1, 2), 1], [(2, 3), 1], [(3, 1), -1], [(4, 2), -1]]

# 参数初始化
w = [0, 0]
b = 0

# 用来记录每次更新过后的w, b
history = []

def update(item):
    """
    随机梯度下降更新参数
    :param item: 参数是分类错误的点
    :return: 无返回值
    """
    # 把w, b, history声明为全局变量
    global w, b, history
    w[0] += item[1] * item[0][0]
    w[1] += item[1] * item[0][1]
    b += item[1]
    history.append([copy.copy(w), b])

def cal(item):
    """
    计算当前模型对样本item的预测值
    :param item: 样本
    :return: 预测值
    """
    res = 0
    for i in range(len(item[0])):
        res += item[0][i] * w[i]
    res += b
    return res

def check():
    """
    检查当前模型是否对所有样本都分类正确
    :return: 如果都正确返回True，否则返回False
    """
    flag = False
    for item in training_set:
        if cal(item) * item[1] <= 0:
            flag = True
            update(item)
    if not flag:
        return True
    else:
        return False

if __name__ == "__main__":
    for i in range(1000):
        if not check():
            pass
        else:
            break
    
    # 绘制动画
    fig = plt.figure()
    ax = plt.axes(xlim=(0, 5), ylim=(-5, 5))
    line, = ax.plot([], [], 'g', lw=2)
    label = ax.text(3, -4, '')
    
    def init():
        line.set_data([], [])
        return line, label
    
    def animate(i):
        # 获取历史数据
        w, b = history[i]
        # 计算直线
        x = range(0, 5)
        y = [(-b - w[0] * xi) / w[1] for xi in x]
        line.set_data(x, y)
        # 设置标签
        label.set_text(f'w={w}, b={b}')
        return line, label
    
    # 创建动画
    anim = animation.FuncAnimation(fig, animate, init_func=init, frames=len(history), interval=1000, blit=True)
    
    # 绘制训练数据点
    for item in training_set:
        if item[1] == 1:
            plt.plot(item[0][0], item[0][1], 'ro')
        else:
            plt.plot(item[0][0], item[0][1], 'bo')
    
    plt.show()